%%

% try different shifts
% include the REST class to soak variance 
   % - in the movement / non movement check
warning('off', 'MATLAB:colon:nonIntegerIndex')
%% segement all the files
FileList = {
   
  %  'mohEEG111.mat'
    'Marion12_2504.mat'
   % 'Tim12_2804.mat'
   %   'Romain12_2904.mat'
    };

%FileList = { 'session2.acq.mat' };
%Fs = 2000; default
Fs = 2000;

window = 3000; % 3000 = 1.5 seconds
shift = 100;
X = []; % contains the trials from all classes
Y = []; % contains the classification from all classes
X_rest = [];
Y_rest = [];

is_pulsar_single_trigger_channel_from_eeglab = false;

for curFileIx = 1:length(FileList)
    
        %Instead of file the data is obtained from EEGLAB 
        if (is_pulsar_single_trigger_channel_from_eeglab)
            
            FileList = {'EEGLAB data'};
            data = EEG.data'; 
            Fs = EEG.srate;
            
            datar = data(:,1:size(data,2)-1); % remove 1 trigger channel at the end

            trigger_channel = data(:, size(data,2));

            classN = length(count_sequence(trigger_channel,2)); % "2" tags each class

            classRest = classN; % the rest class is always the last

            trialsN = length(count_sequence(trigger_channel,1));  % "1" tags each trial
            
        else   
            
            load(['D:\temp\subjects\' FileList{curFileIx}]);
            disp(['Subject file:' FileList{curFileIx}]);
            
            datar = data(:,1:size(data,2)-2); % remove 2 trigger channel at the end

            %trial_channel = data(:,8);
            trial_channel = data(:, size(data,2)-1);

            %class_separation_channel = data(:,9);
            class_separation_channel = data(:, size(data,2));

            classN = length(count_sequence(class_separation_channel,5));  % "5" tags each class
            

            classRest = classN; % the rest class is always the last

            trialsN = length(count_sequence(trial_channel,5));  % "5" tags each trial
            
        end;
        
        disp(['Classes/Movements detected: ' int2str(classN)]);
        disp(['Total trails detected: ' int2str(trialsN)]);
        disp(['Frequency: ' int2str(Fs)]);

        for m = 1:classN
            
            %classes should be supplied to this function
            %select only the trials for this class
            %"best_index" - in 1 shows the postion of a trial, "best_trigger" -
            %again shows the postion but in the range 1 .. n
            if (is_pulsar_single_trigger_channel_from_eeglab)
               %Pulsar GTEC
               [best_index,best_trigger] = generate_trials_postitions_using_single_discrete_channel(m, trigger_channel, 1, 2);
            else
               %Pulsar BIOPAC 
               [best_index,best_trigger] = generate_trials_postitions_using_discrete_channel(m, trial_channel,class_separation_channel);
            end;
            
            % cut signal
            %index = find(diff(trigger)==1); % because we do it earlier now!!!!
            trials_per_class = zeros(size(datar,2),window,length(best_index));
            disp(['Trials: ' int2str(length(best_index)) ' ' FileList{curFileIx}])

            %get the data for these triggers
            for i=1:length(best_index)
                %calculate a range around the trigger
                buffer = datar(best_index(i)+shift:best_index(i)+shift+window-1,:)';
                trials_per_class(:,:,i) = buffer;
                
                buffer_rest_after_mv = datar(best_index(i)+window:best_index(i)+window+3800,:)';
                trials_rest_after_mv(:,:,i) = buffer_rest_after_mv;
            end

            %if (m == classRest)
               %X_rest = cat(3,X_rest,trials_per_class);
               X_rest = cat(3,X_rest,trials_rest_after_mv);
               Y_rest = cat(1,Y_rest,ones(length(best_index),1) * m);
           % else
               X = cat(3,X,trials_per_class); % put all trials together
               Y = cat(1,Y,ones(length(best_index),1) * m);
           % end;
            

        end;
   
    disp('--------------------------------');% next file
end

%% Train on X in CallClass

% the standard MDM

method_mean = 'riemann';
method_dist = 'riemann';

COVtrain = covariances(X);
Ytrain = Y;
labels = unique(Ytrain);
Nclass = length(labels);
CallClass = cell(Nclass,1);

% estimation of center
for i=1:Nclass
    CallClass{i} = mean_covariances(COVtrain(:,:,Ytrain==labels(i)),method_mean);
end

%% Get test data

% we train a potato to detect movements
overlap = window-window/8;

%COV_all_signal=eeg2cov(datar',window,overlap);
COV_movement = covariances(X);
COV_rest = covariances(X_rest);

[C_mv, th] = potato_estimation_iterativ(COV_movement); % train on movement

[C_rest, th_rest] = potato_estimation_iterativ(COV_rest); % train on rest

[COVtest CovTestIndex] = eeg2cov_index(datar',window,overlap);

[~,d_mv] = potato_detection(COVtest,C_mv,th);

[~,d_rest] = potato_detection(COVtest,C_rest,th_rest);

% for i = 1:size(COVtest,3)
%     if (d_rest(i)/th_rest > d_mv(i)/th) % if the REST is stronger than the movement
%         d_mv(i) = 20; % nake the value so high that is not used
%         disp('movement ommitted');
%     end;
% end;


%% detect the good threshold for movement detection
n_trials = (classN-1) * 20;

last = 1000;
best = -1;

for i = 1.40:0.1:8
   
   %skip if the fastest movement is less than ... 400 ms = 800 samples
   minSamples = 600; % 300ms
   minn = window / minSamples; % 1 in d equals 1 window, so we convert minn to windows too
   
   if (count_sequence(d_mv>i,1)) < minn 
       continue;
   end;
   
   if (abs((length(count_sequence(d_mv>i,1)) - n_trials)) < last)
       last = abs(length(count_sequence(d_mv>i,1)) - n_trials)
       best = i
   end;

end;

th = best;

NTesttrial = sum(d_mv<th);
 
COV = zeros(10,10,NTesttrial);

% now we use the potato to extract movements from our signal 
pos = 0;
for i = find(d_mv'<th)%1:length(d)
    
    buffer = datar(CovTestIndex(i):CovTestIndex(i)+window,:);
    cc = covariances(buffer');
    pos = pos + 1;
    COV(:,:,pos) = cc;
 
end;

%% Now we have movements detected and trained classifier in CallClass

COVtest = COV;

d_test = zeros(NTesttrial,Nclass);
for j=1:NTesttrial
    for i=1:Nclass
        d_test(j,i) = distance(COVtest(:,:,j),CallClass{i},method_dist);
    end
end

[~,ix] = min(d_test,[],2);
Ytest = labels(ix);
    
% covariance estimation
disp('done');

